import numpy as np
from rapidocr import RapidOCR

def ocr_to_markdown_text(result, y_thresh=10):
    """
    将 RapidOCR 识别结果转为 Markdown 表格（仅文字内容，按坐标位置排列）
    
    :param result: RapidOCROutput 对象
    :param y_thresh: y坐标的合并阈值（像素），用于判定是否属于同一行
    :return: markdown 文本
    """
    words_with_pos = []

    for txt, box in zip(result.txts, result.boxes):
        # box 是4个点，取平均 y 作为行坐标，平均 x 作为位置
        avg_x = np.mean(box[:, 0])
        avg_y = np.mean(box[:, 1])
        words_with_pos.append((avg_y, avg_x, txt))

    # 按 y 坐标排序
    words_with_pos.sort(key=lambda x: (x[0], x[1]))

    # 分行
    lines = []
    current_line = []
    last_y = None

    for y, x, txt in words_with_pos:
        if last_y is None or abs(y - last_y) <= y_thresh:
            current_line.append((x, txt))
            last_y = y if last_y is None else (last_y + y) / 2  # 平滑行高
        else:
            # 换行
            current_line.sort(key=lambda v: v[0])  # 按 x 排序
            lines.append([t for _, t in current_line])
            current_line = [(x, txt)]
            last_y = y
    if current_line:
        current_line.sort(key=lambda v: v[0])
        lines.append([t for _, t in current_line])

    # 转为 markdown
    md_lines = []
    for line in lines:
        md_lines.append("| " + " | ".join(line) + " |")

    return "\n".join(md_lines)


if __name__ == "__main__":
    engine = RapidOCR()
    img_url = "input_files/1.png"
    result = engine(img_url)

    md_text = ocr_to_markdown_text(result)
    print(md_text)

    # 保存成文件
    with open("ocr_text.md", "w", encoding="utf-8") as f:
        f.write(md_text)
